#include <utils.h>
#include <elf.h>

static bool has_elf = false;
static FILE *elf_fp = NULL;
static int  nr_sym  = 0;
static Elf64_Sym *elf_symtab = NULL;
static char *elf_strtab   = NULL;
static char indent_buf[128];
static uint32_t indent_level  = 0;


bool ftrace_enable() {
  return MUXDEF(CONFIG_FTRACE, has_elf, false);
}

const char *get_indent(uint32_t indent) {
  static int pre_indent = 0;
  indent_buf[pre_indent] = '\t';
  assert(indent < ARRLEN(indent_buf));
  indent_buf[indent] = '\0';
  pre_indent = indent;
  return indent_buf;
}

const char *func_name_off(vaddr_t pc) {
  assert(elf_symtab || elf_strtab);
  for(int i = 0; i < nr_sym; i++) {
    if(ELF64_ST_TYPE(elf_symtab[i].st_info) == STT_FUNC) {
      if(elf_symtab[i].st_value <= pc && elf_symtab[i].st_value + elf_symtab[i].st_size > pc)
        return elf_strtab + elf_symtab[i].st_name;
    }
  }
  return "???";
}

void ftrace_log(vaddr_t pc, const char *f1, bool is_return) {
  if(is_return) {
    indent_level--;
    flog_write("%s" FMT_WORD ": %s  %s\n", get_indent(indent_level), pc, f1, "(ret)");
  }
  else {
    flog_write("%s" FMT_WORD ": %s  %s\n", get_indent(indent_level), pc, f1, "(call)");
    indent_level++;
  }

}

void init_elf(const char *elf_file) {
  if(elf_file == NULL) {
    Log("No elf file: disable ftrace");
    has_elf = false;
    return;
  }
  FILE *fp = fopen(elf_file, "rb");
  Assert(fp, "Can not open '%s'", elf_file);
  elf_fp = fp;
  has_elf = true;
  Log("Elf read: %s", elf_file);

  int ret = 0;
  Elf64_Ehdr ehdr;
  ret = fread(&ehdr, sizeof(ehdr), 1, elf_fp);
  assert(ret == 1);
  Assert(memcmp(ehdr.e_ident, ELFMAG, SELFMAG) == 0, "Not a elf file");

  Elf64_Shdr shdr[ehdr.e_shnum];
  fseek(elf_fp, ehdr.e_shoff, SEEK_SET);
  ret = fread(shdr, ehdr.e_shentsize, ehdr.e_shnum, elf_fp);
  assert(ret == ehdr.e_shnum);

  fseek(elf_fp, shdr[ehdr.e_shstrndx].sh_offset, SEEK_SET);
  Log("shstrndx:%d,%lu", (int)ehdr.e_shentsize, sizeof(Elf64_Shdr));
  char *shstr = (char *) calloc(shdr[ehdr.e_shstrndx].sh_size, sizeof(char));
  ret = fread(shstr, shdr[ehdr.e_shstrndx].sh_size, 1, elf_fp);
  assert(ret == 1);


  for(int i = 0; i < ehdr.e_shnum; i++) {
    if(strcmp(".strtab", shstr + shdr[i].sh_name) == 0) {
      elf_strtab = (char *) calloc(shdr[i].sh_size, sizeof(char));
      fseek(elf_fp, shdr[i].sh_offset, SEEK_SET);
      ret = fread(elf_strtab, shdr[i].sh_size, 1, elf_fp);
      assert(ret == 1);
    }
    else if(shdr[i].sh_type == SHT_SYMTAB) {
      elf_symtab = (Elf64_Sym *) calloc(shdr[i].sh_size, sizeof(char));
      fseek(elf_fp, shdr[i].sh_offset, SEEK_SET);
      ret = fread(elf_symtab, shdr[i].sh_size, 1, elf_fp);
      assert(ret == 1);
      nr_sym = shdr[i].sh_size / sizeof(Elf64_Sym);
    }
  }
  free(shstr);

  Assert(elf_strtab, "ELF strtab not found");
  Assert(elf_symtab, "ELF symtab not found");

  for(int i = 0; i < ARRLEN(indent_buf); i++) {
    indent_buf[i] = '\t';
  }
}